import asyncio
import json
import uuid
from datetime import datetime
from typing import Dict, Set, Optional, Any, Union
import websockets
from websockets.server import WebSocketServerProtocol
from websockets.exceptions import ConnectionClosed, WebSocketException
from fastapi import WebSocket

from models import (
    MessageType, UserMessage, AssistantMessage, StatusUpdate, ErrorMessage,
    HeartbeatMessage, SessionContext, UserInfo, TenantInfo
)
from workflow_manager import workflow_manager
from data_sources import data_source_service
from config import settings
import logging

logger = logging.getLogger(__name__)

class WebSocketServer:
    """WebSocket服务器 - 处理客户端连接和消息路由"""
    
    def __init__(self):
        self.connections: Dict[str, Union[WebSocketServerProtocol, WebSocket]] = {}
        self.session_info: Dict[str, Dict[str, Any]] = {}
        self.heartbeat_tasks: Dict[str, asyncio.Task] = {}
    
    async def register_connection(self, websocket: Union[WebSocketServerProtocol, WebSocket], session_id: str):
        """注册新连接"""
        self.connections[session_id] = websocket
        
        # 获取远程地址
        remote_address = None
        try:
            if hasattr(websocket, 'client'):
                # FastAPI WebSocket
                remote_address = f"{websocket.client.host}:{websocket.client.port}"
            elif hasattr(websocket, 'remote_address'):
                # websockets WebSocketServerProtocol
                remote_address = websocket.remote_address
        except:
            remote_address = "unknown"
        
        self.session_info[session_id] = {
            "connected_at": datetime.now(),
            "last_heartbeat": datetime.now(),
            "remote_address": remote_address
        }
        
        # 启动心跳任务
        self.heartbeat_tasks[session_id] = asyncio.create_task(
            self._heartbeat_task(session_id)
        )
        
        logger.info(f"客户端连接: {session_id} from {remote_address}")
        
        # 发送连接确认
        await self.send_message(session_id, {
            "type": "status",
            "data": {
                "status": "connected",
                "message": "连接成功",
                "session_id": session_id,
                "server_time": datetime.now().isoformat()
            }
        })
    
    async def unregister_connection(self, session_id: str):
        """注销连接"""
        if session_id in self.connections:
            del self.connections[session_id]
        
        if session_id in self.session_info:
            del self.session_info[session_id]
        
        if session_id in self.heartbeat_tasks:
            self.heartbeat_tasks[session_id].cancel()
            del self.heartbeat_tasks[session_id]
        
        logger.info(f"客户端断开: {session_id}")
    
    async def send_message(self, session_id: str, message: Dict[str, Any]):
        """发送消息到指定会话"""
        if session_id not in self.connections:
            logger.warning(f"尝试向不存在的会话发送消息: {session_id}")
            return False
        
        try:
            websocket = self.connections[session_id]
            message_json = json.dumps(message, ensure_ascii=False, default=str)
            
            # 根据WebSocket类型选择发送方法
            if hasattr(websocket, 'send_text'):
                # FastAPI WebSocket
                await websocket.send_text(message_json)
            else:
                # websockets WebSocketServerProtocol
                await websocket.send(message_json)
            return True
        except (ConnectionClosed, WebSocketException) as e:
            logger.warning(f"发送消息失败，连接已断开: {session_id}, {e}")
            await self.unregister_connection(session_id)
            return False
        except Exception as e:
            logger.error(f"发送消息失败: {session_id}, {e}")
            return False
    
    async def broadcast_message(self, message: Dict[str, Any], exclude_sessions: Set[str] = None):
        """广播消息到所有连接"""
        exclude_sessions = exclude_sessions or set()
        
        for session_id in list(self.connections.keys()):
            if session_id not in exclude_sessions:
                await self.send_message(session_id, message)
    
    async def handle_message(self, session_id: str, message_data: Dict[str, Any]):
        """处理收到的消息"""
        try:
            message_type = message_data.get("type")
            data = message_data.get("data", {})
            
            logger.info(f"收到消息: {session_id}, type: {message_type}")
            
            if message_type == MessageType.USER:
                await self._handle_user_message(session_id, data)
            elif message_type == MessageType.HEARTBEAT:
                await self._handle_heartbeat(session_id, data)
            elif message_type == "confirmation_response":
                await self._handle_confirmation_response(session_id, data)
            elif message_type == "get_available_skills":
                await self._handle_get_available_skills(session_id)
            else:
                await self.send_message(session_id, {
                    "type": "error",
                    "data": {
                        "error_code": "UNKNOWN_MESSAGE_TYPE",
                        "error_message": f"未知消息类型: {message_type}"
                    }
                })
        
        except Exception as e:
            logger.error(f"处理消息失败: {session_id}, {e}")
            await self.send_message(session_id, {
                "type": "error",
                "data": {
                    "error_code": "MESSAGE_PROCESSING_ERROR",
                    "error_message": f"消息处理失败: {str(e)}"
                }
            })
    
    async def _handle_user_message(self, session_id: str, data: Dict[str, Any]):
        """处理用户消息"""
        try:
            # 处理字段映射：前端发送'content'，后端期望'message'
            if 'content' in data and 'message' not in data:
                data['message'] = data['content']
            
            # 解析用户消息
            user_message = UserMessage(**data)
            
            # 构建会话上下文
            context = await self._build_session_context(session_id, user_message)
            
            # 更新最后活动时间
            if session_id in self.session_info:
                self.session_info[session_id]["last_activity"] = datetime.now()
            
            # 检查是否需要多Agent协作
            from agents import detect_multi_agent_requirement
            required_agents = detect_multi_agent_requirement(user_message.message)
            
            if required_agents:
                # 启动多Agent协作工作流
                async for result in workflow_manager.create_multi_agent_workflow(
                    session_id, user_message.message, context, required_agents
                ):
                    await self.send_message(session_id, result)
            else:
                # 处理单Agent用户输入
                async for result in workflow_manager.process_user_input(
                    session_id, user_message.message, context
                ):
                    await self.send_message(session_id, result)
        
        except Exception as e:
            logger.error(f"处理用户消息失败: {session_id}, {e}")
            await self.send_message(session_id, {
                "type": "error",
                "data": {
                    "error_code": "USER_MESSAGE_ERROR",
                    "error_message": f"用户消息处理失败: {str(e)}"
                }
            })
    
    async def _handle_confirmation_response(self, session_id: str, data: Dict[str, Any]):
        """处理确认响应"""
        try:
            confirmation_id = data.get("confirmation_id")
            response = data.get("response")
            
            if not confirmation_id or not response:
                await self.send_message(session_id, {
                    "type": "error",
                    "data": {
                        "error_code": "INVALID_CONFIRMATION",
                        "error_message": "确认响应格式无效"
                    }
                })
                return
            
            # 构建会话上下文
            context = await self._build_session_context(session_id)
            
            # 处理确认响应
            async for result in workflow_manager.handle_confirmation_response(
                session_id, confirmation_id, response, context
            ):
                await self.send_message(session_id, result)
        
        except Exception as e:
            logger.error(f"处理确认响应失败: {session_id}, {e}")
            await self.send_message(session_id, {
                "type": "error",
                "data": {
                    "error_code": "CONFIRMATION_RESPONSE_ERROR",
                    "error_message": f"确认响应处理失败: {str(e)}"
                }
            })
    
    # _handle_skill_selection方法已移除，技能选择现在通过user消息处理
    
    async def _handle_get_available_skills(self, session_id: str):
        """处理获取可用技能列表"""
        try:
            from agents import get_available_agents, AGENT_REGISTRY
            
            available_skills = []
            for agent_name in get_available_agents():
                agent_class = AGENT_REGISTRY.get(agent_name)
                if agent_class:
                    agent_instance = agent_class()
                    available_skills.append({
                        "name": agent_name,
                        "description": agent_instance.description,
                        "examples": agent_instance.config.examples
                    })
            
            await self.send_message(session_id, {
                "type": "skills_list",
                "data": {
                    "skills": available_skills
                }
            })
        
        except Exception as e:
            logger.error(f"获取技能列表失败: {session_id}, {e}")
            await self.send_message(session_id, {
                "type": "error",
                "data": {
                    "error_code": "GET_SKILLS_ERROR",
                    "error_message": f"获取技能列表失败: {str(e)}"
                }
            })
    
    async def _handle_heartbeat(self, session_id: str, data: Dict[str, Any]):
        """处理心跳消息"""
        if session_id in self.session_info:
            self.session_info[session_id]["last_heartbeat"] = datetime.now()
        
        # 回复心跳
        await self.send_message(session_id, {
            "type": "heartbeat",
            "data": {
                "timestamp": datetime.now().isoformat(),
                "server_status": "ok"
            }
        })
    
    async def _build_session_context(self, session_id: str, user_message: UserMessage = None) -> SessionContext:
        """构建会话上下文"""
        # 模拟用户信息（实际应用中应该从认证系统获取）
        user_info = UserInfo(
            user_id="demo_user_001",
            name="演示用户",
            department="技术部",
            position="软件工程师",
            email="demo@example.com",
            tenant_id="tenant_001"
        )
        
        # 模拟租户信息
        tenant_info = TenantInfo(
            tenant_id="tenant_001",
            tenant_name="演示公司",
            permissions=["read", "write", "execute"],
            max_users=100,
            features=["knowledge_base", "data_analysis", "form_management"]
        )
        
        # 获取现有工作流状态
        workflow_state = workflow_manager.get_active_workflow(session_id)
        
        context = SessionContext(
            session_id=session_id,
            user_info=user_info,
            tenant_info=tenant_info,
            workflow_state=workflow_state.dict() if workflow_state else None,
            last_activity=datetime.now()
        )
        
        # 如果有用户消息，添加相关信息
        if user_message:
            context.skill_selected = user_message.skill_selected
            context.enable_network = user_message.enable_network
        
        return context
    
    async def _heartbeat_task(self, session_id: str):
        """心跳检测任务"""
        try:
            while session_id in self.connections:
                await asyncio.sleep(settings.WEBSOCKET_HEARTBEAT_INTERVAL)
                
                if session_id not in self.session_info:
                    break
                
                last_heartbeat = self.session_info[session_id].get("last_heartbeat")
                if last_heartbeat:
                    time_since_heartbeat = (datetime.now() - last_heartbeat).total_seconds()
                    if time_since_heartbeat > settings.WEBSOCKET_TIMEOUT:
                        logger.warning(f"心跳超时，断开连接: {session_id}")
                        await self.unregister_connection(session_id)
                        break
        
        except asyncio.CancelledError:
            logger.info(f"心跳任务被取消: {session_id}")
        except Exception as e:
            logger.error(f"心跳任务异常: {session_id}, {e}")
    
    def get_connection_stats(self) -> Dict[str, Any]:
        """获取连接统计信息"""
        current_time = datetime.now()
        
        stats = {
            "total_connections": len(self.connections),
            "active_sessions": len(self.session_info),
            "connections": []
        }
        
        for session_id, info in self.session_info.items():
            connected_duration = (current_time - info["connected_at"]).total_seconds()
            last_activity = info.get("last_activity", info["connected_at"])
            idle_time = (current_time - last_activity).total_seconds()
            
            stats["connections"].append({
                "session_id": session_id,
                "remote_address": str(info["remote_address"]),
                "connected_duration": connected_duration,
                "idle_time": idle_time,
                "is_active": session_id in self.connections
            })
        
        return stats

# 全局WebSocket服务器实例
ws_server = WebSocketServer()

async def websocket_handler(websocket: WebSocketServerProtocol, path: str):
    """WebSocket连接处理器"""
    session_id = str(uuid.uuid4())
    
    try:
        # 注册连接
        await ws_server.register_connection(websocket, session_id)
        
        # 处理消息循环
        async for message in websocket:
            try:
                message_data = json.loads(message)
                await ws_server.handle_message(session_id, message_data)
            except json.JSONDecodeError as e:
                logger.error(f"JSON解析失败: {session_id}, {e}")
                await ws_server.send_message(session_id, {
                    "type": "error",
                    "data": {
                        "error_code": "INVALID_JSON",
                        "error_message": "消息格式错误"
                    }
                })
            except Exception as e:
                logger.error(f"消息处理异常: {session_id}, {e}")
                await ws_server.send_message(session_id, {
                    "type": "error",
                    "data": {
                        "error_code": "MESSAGE_ERROR",
                        "error_message": f"消息处理异常: {str(e)}"
                    }
                })
    
    except ConnectionClosed:
        logger.info(f"客户端正常断开连接: {session_id}")
    except WebSocketException as e:
        logger.warning(f"WebSocket异常: {session_id}, {e}")
    except Exception as e:
        logger.error(f"连接处理异常: {session_id}, {e}")
    finally:
        # 清理连接
        await ws_server.unregister_connection(session_id)

async def start_websocket_server():
    """启动WebSocket服务器"""
    logger.info(f"启动WebSocket服务器: {settings.HOST}:{settings.WEBSOCKET_PORT}")
    
    server = await websockets.serve(
        websocket_handler,
        settings.HOST,
        settings.WEBSOCKET_PORT,
        ping_interval=settings.WEBSOCKET_HEARTBEAT_INTERVAL,
        ping_timeout=settings.WEBSOCKET_TIMEOUT
    )
    
    logger.info(f"WebSocket服务器已启动: ws://{settings.HOST}:{settings.WEBSOCKET_PORT}")
    return server